from aip import AipOcr
import os
import cv2
import numpy as np
import pandas as pd
import requests

spec = ['0301项目实时进展.jpg', '0303项目实时进展.jpg', '0305项目实时进展.jpg', '0307项目实时进展.jpg', '0309项目实时进展.jpg', '0315项目实时进展.jpg',
        '0316项目实时进展.jpg', '0317项目实时进展.jpg', '0318项目实时进展.jpg', '0319项目实时进展.jpg']
spec1 = ['0229项目实时进展.jpg', '0303项目实时进展.jpg', '0305项目实时进展.jpg', '0310项目实时进展.jpg', '0314项目实时进展.jpg', '0318项目实时进展.jpg',
         '0319项目实时进展.jpg']
spec2 = ['0310项目实时进展.jpg', '0314项目实时进展.jpg', '0318项目实时进展.jpg', '0319项目实时进展.jpg']


class ImageTableOCR(object):
    # 初始化
    def __init__(self, ImagePath, X0, X1):
        self.image = cv2.imdecode(np.fromfile(ImagePath, dtype=np.uint8), -1)
        h, w = self.image.shape[:2]
        self.image = cv2.resize(self.image, (2251, int(h * (2251 / w))))[0:int(h * (2251 / w)), X0:X1]
        self.gray = cv2.cvtColor(self.image, cv2.COLOR_BGR2GRAY)
        # 图像二值化
        self.binary = cv2.adaptiveThreshold(~self.gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 35, -5)
        # 进行两次中值滤波

    def HorizontalLineDetect(self):
        # 模板大小3*3
        blur = cv2.medianBlur(self.binary, 5)  # 模板大小3*3
        blur = cv2.medianBlur(blur, 5)
        h, w = self.binary.shape
        # 横向直线列表
        horizontal_lines = []
        for i in range(h - 1):
            # 找到两条记录的分隔线段，以相邻两行的平均像素差大于120为标准
            if abs(np.mean(blur[i, :]) - np.mean(blur[i + 1, :])) > 120:
                # 在图像上绘制线段
                horizontal_lines.append(i)
                # cv2.line(self.image, (0, i), (w, i), (0, 255, 0), 2)
        horizontal_lines = horizontal_lines[25:-15]
        return horizontal_lines

def download_image(dir):
    images = pd.read_csv('image_urls@20200321.csv')
    for name, path in zip(images['name'], images['path']):
        with open('images/' + dir + '/' + name + '.jpg', 'wb') as f:
            image = requests.get(path).content
            print(name, path)
            f.write(image)


def cut_image(dir):
    for root, dirs, files in os.walk(dir):
        for file in files:
            path = os.path.join(dir, file)
            if file in spec:
                X0 = 200
                num = 4020
            else:
                X0 = 310
                num = 4050
            if file in spec1:
                X1 = 2200
            elif file in spec2:
                X1 = cv2.imdecode(np.fromfile(path, dtype=np.uint8), -1).shape[1]
            else:
                X1 = 1997

            im = ImageTableOCR(path, X0, X1)
            hs = im.HorizontalLineDetect()
            y0 = min(hs)
            y1 = max(hs)
            image = im.gray[y0:y1, 0:im.gray.shape[1]]
            if not os.path.exists('pimages/' + file[:4]):
                os.mkdir('pimages/' + file[:4])
            if os.path.exists('pimages/' + '%s.jpg' % file[:4]):
                os.remove('pimages/' + '%s.jpg' % file[:4])
            if y1 - y0 <= num:
                cv2.imwrite('pimages/' + file[:4] + '/1.jpg', image)
            else:
                for i in range(0, y1 - y0, num):
                    tmp = image[i:i + num, 0:image.shape[1]]
                    cv2.imwrite('pimages/' + file[:4] + '/%d.jpg' % int(i / num), tmp)
            print(file)


cut_image('images/images3')
